import os
import pandas as pd
from datetime import datetime, timedelta
from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.object import HistoryRequest
from vnpy.trader.database import get_database
from vnpy.trader.utility import extract_vt_symbol
from vnpy_ctastrategy.backtesting import BacktestingEngine
from vnpy_ctastrategy.backtesting import BacktestingMode
from vnpy_ctastrategy.strategies.atr_rsi_strategy import AtrRsiStrategy  # 导入你的策略类
from vnpy.trader.setting import SETTINGS

# 获取当前工作目录（可能受运行方式影响）
current_working_dir = os.getcwd()
data_dir = current_working_dir + "/data/"

SETTINGS["database.name"] = "mysql"              # 可以根据自己的需求选择数据库，这里使用的是TDengine
SETTINGS["database.database"] = "vnpy"
SETTINGS["database.host"] = "127.0.0.1"
SETTINGS["database.port"] = 3306
SETTINGS["database.user"] = "root"
SETTINGS["database.password"] = "12345678"

class BatchBacktester:
    """批量回测器"""
    
    def __init__(self):
        self.engine = BacktestingEngine()
        self.database = get_database()
        self.results = []
    
    def get_all_stocks(self):
        """获取所有股票列表"""
        # 方法1: 从数据库获取
        # symbols = self.database.get_all_symbols()
        # stock_symbols = [s for s in symbols if s.endswith(('.SS', '.SZ'))]
        
        # 方法2: 从CSV文件获取（如果没有数据库）
        df = pd.read_csv(os.path.join(data_dir,'stock_all_info.csv'))
        stock_symbols = df['代码'].tolist()
        
        return stock_symbols
    
    def prepare_backtest_settings(self, vt_symbol):
        """准备回测设置"""
        # 解析品种代码
        # symbol, exchange = extract_vt_symbol(vt_symbol)
        exchange_str_sx = vt_symbol[0:2]
        symbol = vt_symbol[2:]
         
        exchange = exchange_str_sx.replace("sh", "SSE").replace("sz", "SZSE").replace("bj","BSE")
        vt_symbol = symbol + "." + exchange
        
        # 回测参数
        settings = {
            "vt_symbol": vt_symbol,
            "interval": Interval.DAILY,  # 日线数据
            "start": datetime(2025, 1, 1),
            "end": datetime(2025, 8, 19),
            "rate": 0.0003,  # 手续费率
            "slippage": 0.001,  # 滑点
            "size": 100,  # 合约乘数
            "pricetick": 0.01,  # 价格跳动
            "capital": 50000,  # 初始资金
        }
        
        # 策略参数（可根据需要调整）
        strategy_settings = {
            "atr_length": 22,
            "atr_ma_length": 10,
            "rsi_length": 5,
            "rsi_entry": 16,
            "trailing_percent": 0.8,
            "fixed_size": 100
        }
        
        return settings, strategy_settings
    
    def load_data(self, vt_symbol, start, end, interval):
        """加载历史数据"""
        # symbol, exchange_str = vt_symbol.split(".")
        exchange_str_sx = vt_symbol[0:2]
        symbol = vt_symbol[2:]
         
        exchange = exchange_str_sx.replace("sh", "SSE").replace("sz", "SZSE").replace("bj","BSE")
        exchange = Exchange(exchange)
        # req = HistoryRequest(
        #     symbol=symbol,
        #     exchange=exchange,
        #     interval=interval,
        #     start=start,
        #     end=end
        # )
        
        try:
            data = self.database.load_bar_data( symbol, exchange, interval, start, end)
        
            if len(data) > 50:  # 确保有足够的数据
                return data
            else:
                print(f"数据不足: {vt_symbol}, 数据条数: {len(data)}")
                return None
        except Exception as e:
            print(f"加载数据失败 {vt_symbol}: {e}")
            return None
    
    def run_single_backtest(self, vt_symbol):
        """运行单个股票回测"""
        print(f"开始回测: {vt_symbol}")
        
        # 准备回测设置
        settings, strategy_settings = self.prepare_backtest_settings(vt_symbol)
        
        # 加载数据
        data = self.load_data(
            vt_symbol, 
            settings["start"], 
            settings["end"], 
            settings["interval"]
        )
        
        if not data:
            return None
        
        # 设置回测引擎
        self.engine.set_parameters(
            vt_symbol=settings["vt_symbol"],
            interval=settings["interval"],
            start=settings["start"],
            end=settings["end"],
            rate=settings["rate"],
            slippage=settings["slippage"],
            size=settings["size"],
            pricetick=settings["pricetick"],
            capital=settings["capital"],
            mode=BacktestingMode.BAR
        )
        
        # 添加数据
        self.engine.clear_data()
        self.engine.load_data()
        
        # 添加策略
        self.engine.add_strategy(AtrRsiStrategy, strategy_settings)
        
        # 运行回测
        try:
            self.engine.run_backtesting()
            self.engine.calculate_result()
            statistics =self.engine.calculate_statistics()
            
            # 获取回测结果
            result = {
                "symbol": vt_symbol,
                "total_return": statistics["total_return"],
                "annual_return": statistics["annual_return"],
                "max_drawdown": statistics["max_drawdown"],
                "sharpe_ratio": statistics["sharpe_ratio"],
                "total_net_pnl": statistics["total_net_pnl"],
                "daily_return": statistics["daily_return"],
                "data_points": len(data)
            }
            
            print(f"回测完成: {vt_symbol}")
            return result
            
        except Exception as e:
            print(f"回测失败 {vt_symbol}: {e}")
            return None
    
    def run_batch_backtest(self, max_stocks=None):
        """运行批量回测"""
        print("开始批量回测...")
        
        # 获取所有股票
        all_stocks = self.get_all_stocks()
        if max_stocks:
            all_stocks = all_stocks[:max_stocks]  # 限制回测数量
        
        print(f"共发现 {len(all_stocks)} 只股票")
        
        # 循环回测每只股票
        for i, vt_symbol in enumerate(all_stocks, 1):
            print(f"\n[{i}/{len(all_stocks)}] ", end="")
            
            result = self.run_single_backtest(vt_symbol)
            if result:
                self.results.append(result)
            
            # 每10只股票保存一次结果
            if i % 10 == 0:
                self.save_results()
        
        # 保存最终结果
        self.save_results()
        print("批量回测完成！")
    
    def save_results(self):
        """保存回测结果"""
        if not self.results:
            return
        
        df = pd.DataFrame(self.results)
        
        # 按年化收益排序
        df = df.sort_values("annual_return", ascending=False)
        
        # 保存到CSV
        filename = os.path.join(data_dir, "backtest_results.csv")
        df.to_csv(filename, index=False, encoding="utf-8-sig")
        
        print(f"结果已保存到: {filename}")
        
        # 显示统计信息
        self.show_statistics(df)
    
    def show_statistics(self, df):
        """显示统计信息"""
        print("\n" + "="*50)
        print("回测统计摘要:")
        print(f"有效回测股票数: {len(df)}")
        print(f"平均年化收益: {df['annual_return'].mean():.2%}")
        print(f"最高年化收益: {df['annual_return'].max():.2%}")
        print(f"最低年化收益: {df['annual_return'].min():.2%}")
        print(f"平均最大回撤: {df['max_drawdown'].mean():.2%}")
        print(f"平均夏普比率: {df['sharpe_ratio'].mean():.2f}")
        
        # 显示前10名
        print("\n表现最好的10只股票:")
        top_10 = df.head(10)
        for _, row in top_10.iterrows():
            print(f"{row['symbol']}")

    def generate_report(self):
        """生成详细报告"""
        if not self.results:
            print("没有回测结果")
            return
        
        df = pd.DataFrame(self.results)
        
        # 创建详细报告
        report = {
            "total_stocks": len(df),
            "positive_return": len(df[df["annual_return"] > 0]),
            "negative_return": len(df[df["annual_return"] <= 0]),
            "avg_annual_return": df["annual_return"].mean(),
            "median_annual_return": df["annual_return"].median(),
            "avg_max_drawdown": df["max_drawdown"].mean(),
            "avg_sharpe_ratio": df["sharpe_ratio"].mean()
        }
        
        print("\n详细报告:")
        for key, value in report.items():
            if "return" in key:
                print(f"{key}: {value:.2%}")
            else:
                print(f"{key}: {value:.2f}")

# 使用示例
if __name__ == "__main__":
    # 创建回测器实例
    backtester = BatchBacktester()
    
    # 运行批量回测（限制前50只股票）
    backtester.run_batch_backtest(max_stocks=50)
    
    # 生成详细报告
    backtester.generate_report()
    
    # 也可以单独回测某只股票
    # result = backtester.run_single_backtest("000001.SZ")
